execfile('v4/base.py')

class MaF(Base):
    def __init__(self):
        self.bar_count_default = 50
        self.end_date_str = '20250831'  # 如果需要可以取消注释
        # self.start_date_str = '2025-06-16'  # 如果需要可以取消注释
    
    '''获取日均线'''
    def mark_ma_d(self,df):
        ma_df = df.copy()
        ma_df = ma_df.sort_values(['code', 'date'])
        ma_df['date'] = pd.to_datetime(ma_df['date'])
        
        # 合并分组操作（关键优化点）[7](@ref)
        grouped = ma_df.groupby('code')['close']
        
        # 批量计算所有移动平均（含新增ma8/ma10）[2,4](@ref)
        ma_windows = [3, 5, 8, 10, 15,120,250]
        for window in ma_windows:
            ma_df[f'ma{window}'] = grouped.transform(lambda s: s.rolling(window, min_periods=1).mean())
        
        # 计算乖离率（向量化实现）[2,5](@ref)
        for window in [3, 5, 15,120,250]:
            ma = ma_df[f'ma{window}']
            ma_df[f'bias{window}'] = (ma_df['close'] - ma) / ma * 100
        
        # 计算倾斜率（差分法优化）[6,8](@ref)
        for window in [3, 5, 15,120,250]:
            ma_col = f'ma{window}'
            # 计算当日与前一日MA的差值作为短期斜率
            ma_df[f'slope{window}'] = ma_df.groupby('code')[ma_col].transform(lambda s: s.diff())
        
        # 计算多均线条件（向量化布尔运算）[3](@ref)
        ma_df['d_duo'] = (
            (ma_df['ma3'] > ma_df['ma5']) & 
            (ma_df['ma5'] >= ma_df['ma8']) & 
            (ma_df['ma8'] >= ma_df['ma10'])
        ).astype(int)  # 转为1/0便于后续计算
        ma_df['d_kong'] = (
            (ma_df['ma3'] < ma_df['ma5']) & 
            (ma_df['ma5'] < ma_df['ma8']) & 
            (ma_df['ma8'] < ma_df['ma10'])
        ).astype(int)  # 转为1/0便于后续计算
        # 删除中间列（可选）
        ma_df = self.drop_base_columns(ma_df)
        return ma_df.reset_index(drop=True).round(2)


    '''获取周均线特征'''
    def mark_ma_w(self,df):
        ma_df = df.copy()
        ma_df = ma_df.sort_values(['code', 'date'])
        ma_df['date'] = pd.to_datetime(ma_df['date'])
        
        # 合并分组操作（关键优化点）[7](@ref)
        grouped = ma_df.groupby('code')['close']
        
        # 批量计算所有移动平均（含新增ma8/ma10）[2,4](@ref)
        ma_windows = [3, 5, 8, 10, 15]
        for window in ma_windows:
            ma_df[f'ma{window}w'] = grouped.transform(lambda s: s.rolling(window, min_periods=1).mean())
        
        # 计算乖离率（向量化实现）[2,5](@ref)
        for window in [3, 5, 8]:
            ma = ma_df[f'ma{window}w']
            ma_df[f'bias{window}w'] = (ma_df['close'] - ma) / ma * 100
        
        # 计算倾斜率（差分法优化）[6,8](@ref)
        for window in [3,5,8]:
            ma_col = f'ma{window}w'
            # 计算当日与前一日MA的差值作为短期斜率
            ma_df[f'slope{window}w'] = ma_df.groupby('code')[ma_col].transform(lambda s: s.diff())
        
        # 计算多均线条件（向量化布尔运算）[3](@ref)
        ma_df['w_duo'] = (
            (ma_df['ma3w'] > ma_df['ma5w']) & 
            (ma_df['ma5w'] >= ma_df['ma8w']) & 
            (ma_df['ma8w'] >= ma_df['ma10w'])
        ).astype(int)  # 转为1/0便于后续计算
        # 删除中间列（可选）
        ma_df = self.drop_base_columns(ma_df)
        return ma_df.reset_index(drop=True).round(2)